import os
import sys
import yxn
import util
import time
import json
import socket
import base64
import requests
from log import logger
from datetime import datetime, timedelta

# 开始结束时间
START_TS = int(sys.argv[1])
END_TS   = int(sys.argv[2])

# 地址、采样率等参数
UNISERVER_ADDR = "127.0.0.1"
# web端口、数据端口、用户名密码
UNISERVER_WEB_PORT = 9000
UNISERVER_DATA_PORT = 12201
UNISERVER_USERNAME="admin"
UNISERVER_PASSWORD="admin@123"
# 前端负载均衡采样率
UNISERVER_SAMPLE = 1
# 聚合结果发送协议
UNISERVER_SEND_PROTOCOL = "udp"
# 仪表板ID和查询ID，根据浏览器审查元素获取
UNISERVER_DASHBOARD_ID="66e82380427c647ec4660a7c"
# 主机维度结果集ID
UNISERVER_HOST_QUERY_ID="44f56288-759a-49b9-9854-64ee11ade27a"
# 系统维度结果集ID
UNISERVER_SYSTEM_HOST_QUERY_ID="bbec800e-405f-4dac-abf8-e483a129cf0f"
# 部门维度结果集ID
UNISERVER_DEPARTMENT_HOST_QUERY_ID="ac76e8f7-9b98-46d6-9521-aa2b36cec145"
# 域名维度结果集ID
UNISERVER_DOMAIN_QUERY_ID="4dfcea9e-9307-4640-b5c6-e7972bd99d33"

# 仪表板元素数据
METADATA_FILE = "metadata.json"
# 登录请求session地址
TOKEN_URL = f"http://{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}/api/system/sessions"
# 提交仪表板元数据地址
METADATA_URL = f"http://{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}/api/views/search/metadata"
# 执行查询获取结果地址
EXECUTE_URL = f"http://{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}/api/views/search/{UNISERVER_DASHBOARD_ID}/execute"

# 聚合结果文件保存路径
HOST_ARCH_RESULT_SAVE_PATH = "/uniarch/host"
SYSTEM_HOST_ARCH_RESULT_SAVE_PATH = "/uniarch/system_host"
DEPARTMENT_HOST_ARCH_RESULT_SAVE_PATH = "/uniarch/department_host"

# 模拟登录请求体
TOKEN_DATA = {
  "username": UNISERVER_USERNAME,
  "password": UNISERVER_PASSWORD,
  "host": f"{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}"
}

# 模拟查询仪表板请求体
EXECUTE_DATA = {
  "global_override": {
    "timerange": {
      "type": "absolute",
      "from": None,
      "to": None
    }
  },
  "parameter_bindings": {}
}

# 模拟任何请求的头
REQUEST_HEADERS = {
  "Accept": "application/json",
  "Accept-Language": "zh-CN,zh;q=0.9",
  "Content-Type": "application/json",
  "X-Requested-By": "XMLHttpRequest",
  "X-Requested-With": "XMLHttpRequest"
}

# 处理仪表板每个小部件中的一行
# 'key': ['10.0.8.2'],
# 'values': [{
# 	'key': ['HOST_SESSION_COUNT'],
# 	'value': 29406,
# 	'rollup': True,
# 	'source': 'row-leaf'
# }, {
def process_host_metrics(host_metrics, host_agg_result):
  # 过滤空主机，key中为主机IP
  if len(host_metrics['key']) == 0: return
  # 获取主机和键值对字典
  host = host_metrics['key'][0]
  vals = host_metrics['values']
  # 过滤127.0.0.1主机和非IP地址字符串
  if '127.0.0.1' == host or not util.is_ipaddress(host): return
  # 如果不存在创建主机字典
  if host not in host_agg_result:
    host_agg_result[host] = yxn.create_host_val_dict(host, START_TS, END_TS)
  # 遍历指标键值对字典
  for val_dict in vals:
    key = val_dict['key'][0]
    val = val_dict['value']
    host_agg_result[host][key] = val

# 处理仪表板每个小部件中的一行
# 'key': ['电子警察系统', '10.0.8.2'],
# 'values': [{
# 	'key': ['SYSTEM_CPU_CORE_TOTAL'],
# 	'value': 28,
# 	'rollup': True,
# 	'source': 'row-leaf'
# }, {
def process_system_host_metrics(system_host_metrics, system_host_agg_result):
  if 2 != len(system_host_metrics['key']): return
  # key中为系统名, 主机IP
  system = system_host_metrics['key'][0]
  host   = system_host_metrics['key'][1]
  vals   = system_host_metrics['values']
  # 过滤127.0.0.1主机和非IP地址字符串
  if '127.0.0.1' == host or not util.is_ipaddress(host): return
  # 如果不存在系统创建字典
  if system not in system_host_agg_result:
    system_host_agg_result[system] = yxn.create_system_host_val_dict(system, START_TS, END_TS)
  # 遍历指标键值对字典，累加行中的每一列
  for val_dict in vals:
    key = val_dict['key'][0]
    val = val_dict['value']
    if key not in system_host_agg_result[system]:
      system_host_agg_result[system][key] = 0
    if isinstance(val, str):
      logger.warning(f'MERGE: {system}, {host}, {key}, {val} is string not number ...')
      continue
    system_host_agg_result[system][key] += val

def send_arch_merge_result(agg_result):
  # 发送到GELF HTTP 或者 UDP
  if 'http' == UNISERVER_SEND_PROTOCOL:
    gelf_url = f'http://{UNISERVER_ADDR}:{UNISERVER_DATA_PORT}/gelf'
    with requests.Session() as session:
      for key, val in agg_result.items():
        message = json.dumps(val)
        session.post(gelf_url, data=message, headers={'Content-Type': 'application/json'})
  elif 'udp' == UNISERVER_SEND_PROTOCOL:
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as gelf_udp:
      for key, val in agg_result.items():
        message = json.dumps(val)
        gelf_udp.sendto(message.encode('utf-8'), (UNISERVER_ADDR, UNISERVER_DATA_PORT))

def save_arch_merge_result(dir_path, agg_result):
  # 目录不存在，创建目录
  if not os.path.exists(dir_path): os.makedirs(dir_path)
  # 保存到文件存储
  filename = datetime.fromtimestamp(START_TS).strftime("%Y-%m-%d %H:%M:%S")
  result_file = f'{dir_path}/{filename}.json'
  with open(result_file, 'w', encoding='utf-8') as file:
    json.dump(agg_result, file, ensure_ascii=False, indent=2)
  # 打印保存主机数量
  logger.info(f'MERGE: save {len(agg_result)} records to {result_file} ...')

# 输入主机查询结果，返回主机聚合结果
def arch_host_metrics(host_query_result):
  # 主机聚合结果
  host_agg_result = {}
  # 获取主机维度仪表板查询结果，遍历仪表板的小部件
  for unit_id, unit_result in host_query_result.items():
    # 小部件包括多个主机，每个主机有很多指标
    host_metrics_list = unit_result['rows']
    # 遍历每个主机的指标字典
    for host_metrics in host_metrics_list:
      process_host_metrics(host_metrics, host_agg_result)
  
  # 计算主机比率平均值
  for host, val_dict in host_agg_result.items():
    if 'HOST_CPU_CORE' in val_dict and not isinstance(val_dict['HOST_CPU_CORE'], str) and not isinstance(val_dict['HOST_MEM_AVG'], str):
      val_dict['HOST_CPU_CORE'] = int(val_dict['HOST_CPU_CORE'])
      val_dict['HOST_MEM_AVG'] = int(val_dict['HOST_MEM_AVG'])
      val_dict['HOST_DISK_AVG'] = int(val_dict['HOST_DISK_AVG'])
      yxn.calc_host_avg_ratio(val_dict)
      yxn.calc_host_cloud_score(val_dict)
      for key, val in val_dict.items():
        if key in yxn.REVERT_SAMPLE_KEYS:
          val_dict[key] *= UNISERVER_SAMPLE
      val_dict['timestamp'] = START_TS

  return host_agg_result

# 输入系统查询结果，返回系统聚合结果
def arch_system_host_metrics(system_host_query_result):
  # 系统聚合结果
  system_host_agg_result = {}
  # 获取系统维度仪表板查询结果，遍历仪表板的小部件
  for unit_id, unit_result in system_host_query_result.items():
    # 小部件包括多个系统，每个系统有很多指标
    system_host_metrics_list = unit_result['rows']
    # 遍历每个系统的指标字典
    for system_host_metrics in system_host_metrics_list:
      process_system_host_metrics(system_host_metrics, system_host_agg_result)
  return system_host_agg_result

# 模拟登录查询仪表板
def query_dashboard():
  # 读取仪表板查询JSON
  metadata_data = None
  with open(METADATA_FILE, "r") as f: metadata_data = json.load(f)
  
  # 设置查询时间
  EXECUTE_DATA['global_override']['timerange']['from'] = datetime.utcfromtimestamp(START_TS).isoformat() + "Z"
  EXECUTE_DATA['global_override']['timerange']['to'] = datetime.utcfromtimestamp(END_TS).isoformat() + "Z"
  logger.info(EXECUTE_DATA)
  
  # 模拟登录获取session
  logger.info(f'get token from {TOKEN_URL} ...')
  response = requests.post(TOKEN_URL, headers=REQUEST_HEADERS, json=TOKEN_DATA, verify=False)
  logger.info(f'token response {response.text} ...')
  session_id = json.loads(response.text).get('session_id', None)
  
  # 根据session_id生成token令牌
  uniserver_token = base64.b64encode(f'{session_id}:session'.encode('utf-8')).decode('utf-8')
  REQUEST_HEADERS["Authorization"] = f"Basic {uniserver_token}"
  logger.info(f'token is: {uniserver_token} ...')
  
  # 查询仪表板
  logger.info(f'query {START_TS} to {END_TS} ...')
  requests.post(METADATA_URL, headers=REQUEST_HEADERS, json=metadata_data, verify=False)
  response = requests.post(EXECUTE_URL, headers=REQUEST_HEADERS, json=EXECUTE_DATA, verify=False)
  
  # 解析查询结果
  query_dashboard_result = None
  try:
    query_dashboard_result = json.loads(response.text)
    if "results" not in query_dashboard_result:
      logger.warning(f'no query results in {response.text}')
  except:
    logger.warning(f'failed to load {response.text}')

  return query_dashboard_result

# 查询仪表板
query_dashboard_result = query_dashboard()

# 处理主机仪表板查询结果
host_query_result = query_dashboard_result["results"][UNISERVER_HOST_QUERY_ID]["search_types"]
host_agg_result = arch_host_metrics(host_query_result)
# 发送保存主机聚合结果
save_arch_merge_result(HOST_ARCH_RESULT_SAVE_PATH, host_agg_result)
send_arch_merge_result(host_agg_result)

# 处理系统仪表板查询结果
system_host_query_result = query_dashboard_result["results"][UNISERVER_SYSTEM_HOST_QUERY_ID]["search_types"]
system_host_agg_result = arch_system_host_metrics(system_host_query_result)
# 发送保存系统聚合结果
save_arch_merge_result(SYSTEM_HOST_ARCH_RESULT_SAVE_PATH, system_host_agg_result)
send_arch_merge_result(system_host_agg_result)

# 处理部门仪表板查询结果
department_host_query_result = query_dashboard_result["results"][UNISERVER_DEPARTMENT_HOST_QUERY_ID]["search_types"]
department_host_agg_result = arch_system_host_metrics(department_host_query_result)
# 发送保存部门聚合结果
save_arch_merge_result(DEPARTMENT_HOST_ARCH_RESULT_SAVE_PATH, department_host_agg_result)
send_arch_merge_result(department_host_agg_result)
